import serial
import time
import os
import sys
import random
import re
import subprocess
import threading


# 定义常量
DEBUG_PORT = "COM4"
AT_PORT = "COM32"
BAUDRATE = 115200
# 定义变量
at_serial = None
debug_serial = None
# 记录运行次数
runTimes = 0
# log文件
DEBUG_LOG = ""
ADB_LOG = ""
AT_LOG = ""


# 创建日志文件
def createLogFile():
    global DEBUG_LOG, ADB_LOG, AT_LOG
    timestamp = time.strftime("%Y-%m-%d_%H_%M_%S", time.localtime())
    DEBUG_LOG = f"DEBUG_LOG_{timestamp}.log"
    ADB_LOG = f"ADB_LOG_{timestamp}.log"
    AT_LOG = f"AT_LOG_{timestamp}.log"

    with open(DEBUG_LOG, "w", encoding="utf-8") as file:
        file.write("DEBUG 日志开始记录...\n")

    with open(ADB_LOG, "w", encoding="utf-8") as file:
        file.write("ADB 日志开始记录...\n")

    with open(AT_LOG, "w", encoding="utf-8") as file:
        file.write("AT 日志开始记录...\n")


# 写日志
def write_log(log, logFile):
    with open(logFile, "a", encoding="utf-8") as file_object:
        file_object.write(
            str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
            + "\r\n"
            + log
            + "\r\n"
        )
        now_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        print("[" + now_time + "]" + " " + log)


# Send at instruction
def at_write(command, timeout, at_port_opened, at_result, RDY=0, CPIN=0):
    at_command = command.encode("UTF-8")
    at_port_opened.write(at_command)


# 读取at口返回
def at_read(serial_port, logFile=""):
    rec_str = ""
    while True:
        time.sleep(1)
        lenth = serial_port.inWaiting()
        if lenth > 0:
            rec_str = serial_port.read(size=lenth)
            try:
                rec_str = re.sub("\\x1b.*?m", "", rec_str.decode("UTF-8"), 0, re.I)
                if logFile != "":
                    with open(logFile, "a") as file_object:
                        rec_str = (
                            str(time.time())
                            + ":"
                            + str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
                            + "\r\n"
                            + rec_str
                        )
                        file_object.write(
                            str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
                            + "\r\n"
                            + rec_str
                            + "\r\n"
                        )
                return rec_str
            except:
                return ""
        else:
            break
    return rec_str


# 登录设备
def loginDevice():
    global DEBUG_LOG
    # 确认登录失败的条件，或者，不可能失败
    while True:
        #   print("等待登录... ... ")
        write_log("等待登录... ... ", DEBUG_LOG)
        time.sleep(3)
        at_write("\r\n", 3, debug_serial, "OK")
        result = at_read(debug_serial)

        with open(DEBUG_LOG, "a") as file_object:
            file_object.write(result)

        if "login:" in result:
            at_write("root\r\n", 3, debug_serial, "OK")
            write_log("账号输入成功", DEBUG_LOG)
            time.sleep(2)
            at_write("oelinux123\r\n", 3, debug_serial, "OK")
            write_log("密码输入成功", DEBUG_LOG)
        elif "~ #" in result:
            write_log("登录成功", DEBUG_LOG)
            break
        else:
            continue


# 检查debug口通信状态
def check_port_status():
    global debug_serial
    if debug_serial == None:
        try:
            debug_serial = serial.Serial(DEBUG_PORT, BAUDRATE)
        except:
            print("端口连接失败，请检查端口")
            sys.exit()


# 写入运行次数到日志
def writeRuntimesToLog():
    global runTimes
    with open(DEBUG_LOG, "a+") as file_object:
        file_object.write(
            "\r\n{} runtimes: {} {}\r\n".format("<" * 20, runTimes, ">" * 20)
        )
    with open(AT_LOG, "a+") as file_object:
        file_object.write(
            "\r\n{} runtimes: {} {}\r\n".format("<" * 20, runTimes, ">" * 20)
        )
    with open(ADB_LOG, "a+") as file_object:
        file_object.write(
            "\r\n{} runtimes: {} {}\r\n".format("<" * 20, runTimes, ">" * 20)
        )
    print("\r\n{} runtimes: {} {}\r\n".format("<" * 20, runTimes, ">" * 20))
    runTimes += 1


# 发送AT指令
def send_at_command(command):
    # 打开串口
    ser = serial.Serial(AT_PORT, BAUDRATE)
    ser.write((command + "\r\n").encode())
    ser.write(("\r\n").encode())
    time.sleep(2)
    response = ser.read_all().decode()
    ser.close()
    # write_log(f"at指令{command}执行结果:{response}",AT_LOG)
    with open(AT_LOG, "a", encoding="utf-8") as file_object:
        file_object.write(
            str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
            + "\r\n"
            + response
            + "\r\n"
        )
    return response


# 发送adb命令
def send_adb_command(command):
    try:
        # 开启adb root
        subprocess.run(["adb", "root"], capture_output=True, text=True)
        # 使用 subprocess 模块执行 adb shell 命令
        result = subprocess.run(
            ["adb", "shell", command], capture_output=True, text=True
        )
        with open(ADB_LOG, "a", encoding="utf-8") as file_object:
            file_object.write(
                str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
                + "\r\n"
                + result.stdout
                + "\r\n"
            )
        # 返回命令的输出结果
        return result.stdout
    except Exception as e:
        return str(e)


# 写入运行次数到日志
def writeRuntimesToLog():
    global runTimes
    with open(DEBUG_LOG, "a+") as file_object:
        file_object.write(
            "\r\n{} runtimes: {} {}\r\n".format("<" * 20, runTimes, ">" * 20)
        )
    print("\r\n{} runtimes: {} {}\r\n".format("<" * 20, runTimes, ">" * 20))
    runTimes += 1


# 擦除分区
def erase_partition(partition):
    at_write(f"mtdimg --erase {partition}\r\n", 3, debug_serial, "OK")
    result = at_read(debug_serial, DEBUG_LOG)
    if "mtd: erase block" in result:
        write_log(f"{partition}分区擦除成功\n", DEBUG_LOG)
    else:
        write_log(f"{partition}分区擦除失败\n", DEBUG_LOG)
    # print(result)
    return result


# 查询还原状态
def query_restore_status():
    at_write("fotainfo --get-all\r\n", 3, debug_serial, "OK")
    result = at_read(debug_serial, DEBUG_LOG)
    # print(result)
    
    return result


# 校验还原状态
def check_restore_status():
    for i in range(0, 61):
        result = query_restore_status()
        time.sleep(2)
        if "switch position:NULL" in result and "FOTA STAT:SUCCEED (0)" in result:
            write_log("状态已就绪可进行擦除操作\n", DEBUG_LOG)
            break
        elif i == 60:
            write_log("状态异常，请检查\n", DEBUG_LOG)
            sys.exit()


def restart_module():
    at_write("reboot\r\n", 3, debug_serial, "OK")
    write_log("等待重启... ... \n", DEBUG_LOG)
    start_time = time.time()  # 记录开始重启的时间戳
    # print(start_time)
    for i in range(0, 61):
        result = at_read(debug_serial, DEBUG_LOG)
        time.sleep(2)

        if "login:" in result:
            end_time = time.time()  # 记录成功的时间戳
            duration = end_time - start_time  # 计算重启耗时
            write_log(f"reboot重启成功，耗时 {duration:.2f} 秒\n", DEBUG_LOG)
            break

        elif i == 60:
            write_log("reboot重启失败\n", DEBUG_LOG)
            sys.exit()

        time.sleep(2)
    return end_time




# 查询对应分区状态
def query_partition_status():
    write_log("查询分区状态... ... \n", AT_LOG)
    result = send_at_command('AT+QNAND="RecoveryCnt"')
    # print(result)
    return result


# 校验对应分区状态
def parse_qnand_response(response):
    data_dict = {}
    lines = response.strip().split("\n")
    for line in lines:
        line = line.strip()
        if ":" in line and not line.startswith("+QNAND"):
            try:
                key, value = line.replace('"', "").split(":")
                data_dict[key.strip()] = int(value.strip())
            except ValueError:
                data_dict[key.strip()] = value.strip()
    return data_dict


def check_partition_status(data_prev, data_curr, partition):
    # 提取上一次数据的还原次数
    count_prev = data_prev.get(partition, -1)
    # 提取当前数据的还原次数
    count_curr = data_curr.get(partition, -1)
    if count_prev == -1 or count_curr == -1:
        write_log(f"分区 {partition} 不存在", DEBUG_LOG)
    # 校验是否增加 1
    if count_curr == count_prev + 1:
        write_log(
            f"分区 {partition} 校验通过,上次值: {count_prev}，当前值: {count_curr}",
            AT_LOG,
        )
        # return f"分区 {partition} 校验通过"
    else:
        write_log(
            f"分区 {partition} 校验失败, 上次值: {count_prev}，当前值: {count_curr}",
            AT_LOG,
        )
        # return f"分区 {partition} 校验失败，上次值: {count_prev}，当前值: {count_curr}"


def operate_partition(partition):
    """
    对指定分区进行擦除、重启及状态检查操作。
    """
    time.sleep(5)
    check_restore_status()
    data = query_partition_status()
    data_prev = parse_qnand_response(data)
    erase_partition(partition)
    # 重启模块
    restart_module()
    # 登录设备
    loginDevice()
    time.sleep(5)
    check_restore_status()
    data = query_partition_status()
    data_curr = parse_qnand_response(data)
    check_partition_status(data_prev, data_curr, partition)


def main():
    # 初始化日志文件
    createLogFile()
    # 检查并打开调试端口
    check_port_status()
    # 登录设备
    loginDevice()
    """ 
    定义需要操作的分区及其 AB 状态
    这里的字典信息用于调试，实际使用时，需要根据实际情况修改，
    增加一条或多条分区及其 AB 状态
    （True/False）表示是否有 AB 分区，true 表示有，false 表示没有
    """
    partitions = {
        # "modem": True,
        # "boot": True,
        # "rawdata": False,
        # "usrdata": False,
        # "abl": True,
        "tz" : True,
    }
    while True:
    # 写入运行次数到日志
        writeRuntimesToLog()
    # 遍历分区，动态处理是否需要 AB 分区
        for partition, has_ab in partitions.items():
            suffixes = ["", "_b"] if has_ab else [""]
            for suffix in suffixes:
                current_partition = f"{partition}{suffix}"
                operate_partition(current_partition)



if __name__ == "__main__":
    main()
